#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <errno.h>

#define BUFFER_SIZE 4096
#define MAX_CLIENTS 10

typedef struct {
    int tcp_socket;
    struct sockaddr_in client_addr;
} client_info_t;

typedef struct {
    int tcp_listen_fd;
    int udp_fd;
    client_info_t clients[MAX_CLIENTS];
    int client_count;
    struct sockaddr_in udp_target;
    int udp_listen_port;
} forward_ctx_t;

// 初始化上下文
int init_forward_ctx(forward_ctx_t *ctx, int tcp_port, const char *udp_ip, int udp_target_port, int udp_listen_port) {
    memset(ctx, 0, sizeof(forward_ctx_t));
    
    // 创建TCP监听socket
    ctx->tcp_listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (ctx->tcp_listen_fd < 0) {
        perror("TCP socket creation failed");
        return -1;
    }
    
    // 设置TCP socket选项
    int opt = 1;
    if (setsockopt(ctx->tcp_listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("setsockopt failed");
        close(ctx->tcp_listen_fd);
        return -1;
    }
    
    // 绑定TCP socket
    struct sockaddr_in tcp_addr;
    memset(&tcp_addr, 0, sizeof(tcp_addr));
    tcp_addr.sin_family = AF_INET;
    tcp_addr.sin_addr.s_addr = INADDR_ANY;
    tcp_addr.sin_port = htons(tcp_port);
    
    if (bind(ctx->tcp_listen_fd, (struct sockaddr*)&tcp_addr, sizeof(tcp_addr)) < 0) {
        perror("TCP bind failed");
        close(ctx->tcp_listen_fd);
        return -1;
    }
    
    // 监听TCP socket
    if (listen(ctx->tcp_listen_fd, 5) < 0) {
        perror("TCP listen failed");
        close(ctx->tcp_listen_fd);
        return -1;
    }
    
    // 创建UDP socket
    ctx->udp_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (ctx->udp_fd < 0) {
        perror("UDP socket creation failed");
        close(ctx->tcp_listen_fd);
        return -1;
    }
    
    // 绑定UDP socket到监听端口
    struct sockaddr_in udp_listen_addr;
    memset(&udp_listen_addr, 0, sizeof(udp_listen_addr));
    udp_listen_addr.sin_family = AF_INET;
    udp_listen_addr.sin_addr.s_addr = INADDR_ANY;
    udp_listen_addr.sin_port = htons(udp_listen_port);
    
    if (bind(ctx->udp_fd, (struct sockaddr*)&udp_listen_addr, sizeof(udp_listen_addr)) < 0) {
        perror("UDP bind failed");
        close(ctx->tcp_listen_fd);
        close(ctx->udp_fd);
        return -1;
    }
    
    // 设置UDP目标地址
    memset(&ctx->udp_target, 0, sizeof(ctx->udp_target));
    ctx->udp_target.sin_family = AF_INET;
    ctx->udp_target.sin_port = htons(udp_target_port);
    
    if (inet_pton(AF_INET, udp_ip, &ctx->udp_target.sin_addr) <= 0) {
        perror("Invalid UDP target IP address");
        close(ctx->tcp_listen_fd);
        close(ctx->udp_fd);
        return -1;
    }
    
    ctx->udp_listen_port = udp_listen_port;
    printf("Forward service initialized:\n");
    printf("  TCP listening on port: %d\n", tcp_port);
    printf("  UDP target: %s:%d\n", udp_ip, udp_target_port);
    printf("  UDP listening on port: %d\n", udp_listen_port);
    
    return 0;
}

// 添加TCP客户端
int add_tcp_client(forward_ctx_t *ctx, int client_fd, struct sockaddr_in *client_addr) {
    if (ctx->client_count >= MAX_CLIENTS) {
        printf("Max clients reached, rejecting new connection\n");
        close(client_fd);
        return -1;
    }
    
    ctx->clients[ctx->client_count].tcp_socket = client_fd;
    memcpy(&ctx->clients[ctx->client_count].client_addr, client_addr, sizeof(*client_addr));
    ctx->client_count++;
    
    char client_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &client_addr->sin_addr, client_ip, INET_ADDRSTRLEN);
    printf("New TCP client connected: %s:%d (total clients: %d)\n", 
           client_ip, ntohs(client_addr->sin_port), ctx->client_count);
    
    return 0;
}

// 移除TCP客户端
void remove_tcp_client(forward_ctx_t *ctx, int index) {
    if (index < 0 || index >= ctx->client_count) {
        return;
    }
    
    char client_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &ctx->clients[index].client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
    printf("TCP client disconnected: %s:%d\n", 
           client_ip, ntohs(ctx->clients[index].client_addr.sin_port));
    
    close(ctx->clients[index].tcp_socket);
    
    // 移动数组元素填补空缺
    for (int i = index; i < ctx->client_count - 1; i++) {
        ctx->clients[i] = ctx->clients[i + 1];
    }
    
    ctx->client_count--;
}

// 处理TCP到UDP的数据转发
void forward_tcp_to_udp(forward_ctx_t *ctx, int client_index) {
    char buffer[BUFFER_SIZE];
    ssize_t bytes_read;
    
    bytes_read = read(ctx->clients[client_index].tcp_socket, buffer, BUFFER_SIZE);
    
    if (bytes_read > 0) {
        // 通过UDP发送数据到目标
        ssize_t bytes_sent = sendto(ctx->udp_fd, buffer, bytes_read, 0,
                                   (struct sockaddr*)&ctx->udp_target, sizeof(ctx->udp_target));
        
        if (bytes_sent < 0) {
            perror("UDP send failed");
        } else {
            printf("Forwarded %zd bytes from TCP client to UDP target\n", bytes_sent);
        }
    } else if (bytes_read == 0) {
        // TCP连接关闭
        remove_tcp_client(ctx, client_index);
    } else {
        perror("TCP read failed");
        remove_tcp_client(ctx, client_index);
    }
}

// 处理UDP到TCP的数据转发
void forward_udp_to_tcp(forward_ctx_t *ctx) {
    char buffer[BUFFER_SIZE];
    struct sockaddr_in udp_src_addr;
    socklen_t addr_len = sizeof(udp_src_addr);
    
    ssize_t bytes_read = recvfrom(ctx->udp_fd, buffer, BUFFER_SIZE, 0,
                                 (struct sockaddr*)&udp_src_addr, &addr_len);
    
    if (bytes_read > 0) {
        // 检查UDP数据包是否来自我们期望的端口（可选）
        // 这里我们转发给所有TCP客户端
        for (int i = 0; i < ctx->client_count; i++) {
            ssize_t bytes_sent = write(ctx->clients[i].tcp_socket, buffer, bytes_read);
            if (bytes_sent < 0) {
                perror("TCP send failed");
                // 如果发送失败，移除该客户端
                remove_tcp_client(ctx, i);
                i--; // 因为数组元素移动了，需要调整索引
            } else {
                char src_ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, &udp_src_addr.sin_addr, src_ip, INET_ADDRSTRLEN);
                printf("Forwarded %zd bytes from UDP %s:%d to TCP client\n", 
                       bytes_sent, src_ip, ntohs(udp_src_addr.sin_port));
            }
        }
    } else if (bytes_read < 0) {
        perror("UDP receive failed");
    }
}

// 主转发循环
void run_forward_loop(forward_ctx_t *ctx) {
    fd_set read_fds;
    int max_fd;
    
    printf("Starting forward loop...\n");
    
    while (1) {
        FD_ZERO(&read_fds);
        
        // 添加TCP监听socket
        FD_SET(ctx->tcp_listen_fd, &read_fds);
        max_fd = ctx->tcp_listen_fd;
        
        // 添加UDP socket
        FD_SET(ctx->udp_fd, &read_fds);
        if (ctx->udp_fd > max_fd) {
            max_fd = ctx->udp_fd;
        }
        
        // 添加所有TCP客户端socket
        for (int i = 0; i < ctx->client_count; i++) {
            FD_SET(ctx->clients[i].tcp_socket, &read_fds);
            if (ctx->clients[i].tcp_socket > max_fd) {
                max_fd = ctx->clients[i].tcp_socket;
            }
        }
        
        // 使用select等待事件
        int activity = select(max_fd + 1, &read_fds, NULL, NULL, NULL);
        
        if (activity < 0) {
            perror("select failed");
            break;
        }
        
        // 检查新的TCP连接
        if (FD_ISSET(ctx->tcp_listen_fd, &read_fds)) {
            struct sockaddr_in client_addr;
            socklen_t client_len = sizeof(client_addr);
            
            int client_fd = accept(ctx->tcp_listen_fd, (struct sockaddr*)&client_addr, &client_len);
            if (client_fd >= 0) {
                add_tcp_client(ctx, client_fd, &client_addr);
            } else {
                perror("TCP accept failed");
            }
        }
        
        // 检查UDP数据
        if (FD_ISSET(ctx->udp_fd, &read_fds)) {
            forward_udp_to_tcp(ctx);
        }
        
        // 检查TCP客户端数据
        for (int i = 0; i < ctx->client_count; i++) {
            if (FD_ISSET(ctx->clients[i].tcp_socket, &read_fds)) {
                forward_tcp_to_udp(ctx, i);
            }
        }
    }
}

// 清理资源
void cleanup_forward_ctx(forward_ctx_t *ctx) {
    close(ctx->tcp_listen_fd);
    close(ctx->udp_fd);
    
    for (int i = 0; i < ctx->client_count; i++) {
        close(ctx->clients[i].tcp_socket);
    }
    
    printf("Forward service stopped\n");
}

// 主函数 - 接口函数
int start_forward_service(int porta, const char *ipu, int portb, int portc) {
    forward_ctx_t ctx;
    
    if (init_forward_ctx(&ctx, porta, ipu, portb, portc) < 0) {
        fprintf(stderr, "Failed to initialize forward service\n");
        return -1;
    }
    
    // 运行主循环
    run_forward_loop(&ctx);
    
    // 清理资源
    cleanup_forward_ctx(&ctx);
    
    return 0;
}

// 使用示例
int main() {
    // 示例用法：
    // TCP监听端口 8080
    // UDP目标地址 192.168.1.100:9090  
    // UDP监听端口 7070
    return start_forward_service(8080, "192.168.1.100", 9090, 7070);
}
